﻿#pragma once

#include "coroutine/coroutine.h"
#include "coroutine/detail/timer.h"
#include "object_pool.h"
#include "rpc_defines.h"
#include "rpc_root.h"
#include "rpc_singleton.h"
#include "rpc_transport.h"
#include <ctime>
#include <memory>
#include <unordered_map>

namespace kratos {
namespace service {
class ServiceContext;
}
} // namespace kratos

namespace rpc {

class ProxyCall;
class Rpc;
class RpcImpl;
class ProxyManager;
class ProxyCallManagerImpl;

// Returns ProxyManager
extern ProxyManager *getProxyManager(Rpc *rpc);

// Service Proxy interface
class Proxy {
public:
  virtual ~Proxy() {}
  // Get method name like ServiceName::methodName
  // @param methodID method sequence index start from 1
  virtual const char *getSignature(MethodID methodID) const = 0;
  // Returns remote service's UUID
  virtual ServiceUUID getServiceUUID() = 0;
  // Is oneway method?
  virtual bool isOneway(MethodID methodID) = 0;
  virtual GlobalIndex getGlobalIndex() = 0;
  // Setup Rpc
  virtual void setRpc(Rpc *rpc) = 0;
  // Returns Rpc
  virtual Rpc *getRpc() = 0;
  // Setup GlobalIndex
  virtual void setGlobalIndex(GlobalIndex globalIndex) = 0;
  // Is the channel of proxy is connected
  virtual bool isConnected() = 0;
};

// TODO 定期清理无效的proxy weak_ptr

// The proxy for user defined remote service API
class ProxyImpl : public virtual Proxy {
  TransportPtr transport_{nullptr}; ///< User implemented transportation
  ServiceID targetID_{
      INVALID_SERVICE_ID};       ///< Attached remote service instance ID
  ProxyID id_{INVALID_PROXY_ID}; ///< Proxy instance ID
  bool attached_{false};         ///< Need attach to remote service instance?
  GlobalIndex globalIndex_{INVALID_GLOBAL_INDEX}; ///< Global index
  RpcImpl *rpc_{nullptr};                         ///< Rpc instance

public:
  // dtor
  virtual ~ProxyImpl();
  // Returns @see Transport pointer
  TransportPtr getTransport();
  // Returns flag to attach remote service or not
  bool isAttached();
  // Returns proxy instance ID
  ProxyID getID();
  // Returns target(attached) service instance ID
  rpc::ServiceID getTarget();
  // Return global index
  virtual GlobalIndex getGlobalIndex() override;
  // Returns Rpc
  virtual Rpc *getRpc() override;
  virtual void setRpc(Rpc *rpc) override;
  // Setup global index
  virtual void setGlobalIndex(GlobalIndex globalIndex) override;
  // Is the channel of proxy is connected
  virtual bool isConnected() override;
  // Setup attached remote service instance ID
  void setTarget(ServiceID serviceID);
  // Setup remote service's Transport interface
  void setTransport(TransportPtr &transport);
  // Open attach flag
  void attach();
  // Close attach flag
  void detach();
  // Setup proxy instance ID
  void setID(ProxyID id);
};

// ProxyManager interface
class ProxyManager {
public:
  virtual ~ProxyManager() {}
  // Add to manager
  virtual void add(ProxyPtr &proxy) = 0;
  // Destory proxy
  // @param Proxy instance ID
  virtual void destroy(ProxyID id) = 0;
  // Find proxy
  // @param Proxy instance ID
  virtual ProxyPtr get(ProxyID id) = 0;
  // Cleanup
  virtual void cleanup() = 0;
};

// The manager to manage proxy
class ProxyManagerImpl : public ProxyManager {
  using ProxyMap = std::unordered_map<ProxyID, ProxyWeakPtr>;
  ProxyMap proxyMap_;     ///< Proxy map
  RpcImpl *rpc_{nullptr}; ///< Rpc instance

public:
  ProxyManagerImpl(RpcImpl *rpc);
  // dtor
  virtual ~ProxyManagerImpl();
  // Add to manager
  virtual void add(ProxyPtr &proxy) override;
  // Destory proxy
  // @param Proxy instance ID
  virtual void destroy(ProxyID id) override;
  // Find proxy
  // @param Proxy instance ID
  virtual ProxyPtr get(ProxyID id) override;
  // Cleanup
  virtual void cleanup() override;
};

// Base interface of proxy factory
class ProxyFactory {
public:
  // dtor
  virtual ~ProxyFactory() {}
  // Create a proxy instance
  virtual ProxyPtr create(Rpc *rpc) = 0;
  // Get proxy class name
  virtual const char *name() = 0;
};

// Template class to create user service's proxy
template <typename T> class ProxyFactoryT : public ProxyFactory {
public:
  // Create a proxy instance
  virtual ProxyPtr create(Rpc *rpc) override {
    auto proxy = std::dynamic_pointer_cast<Proxy>(rpc->make_shared_ptr<T>());
    proxy->setRpc(rpc);
    return proxy;
  }
  // Get proxy class name
  virtual const char *name() override { return typeid(T).name(); }
};

class ProxyCreator {
public:
  virtual ~ProxyCreator() {}
  // Create a proxy
  // @param serviceUUID service UUID
  // @param trans data stream
  // @param attached attach to a service instance or not
  // @param rpc Rpc
  // @return Proxy pointer
  virtual ProxyPtr getService(ServiceUUID serviceUUID, TransportPtr &trans,
                              bool attached, Rpc *rpc) = 0;
  // Create a proxy
  // @param className The name of class
  // @param trans data stream
  // @param attached attach to a service instance or not
  // @param rpc Rpc
  // @return Proxy pointer
  virtual ProxyPtr getService(const std::string &className, TransportPtr &trans,
                              bool attached, Rpc *rpc) = 0;
  // Registers a service factory
  // @param serviceUUID service UUID
  // @param factory service factory
  virtual void addFactory(ServiceUUID serviceUUID, ProxyFactory *factory) = 0;
  // Make a new proxy instance ID
  virtual ProxyID createProxyID() = 0;
};

// Proxy creator
class ProxyCreatorImpl : public ProxyCreator {
  using ProxyFactoryMap =
      std::unordered_map<ServiceUUID, std::unique_ptr<ProxyFactory>>;
  using ProxyFactoryNameIDMap = std::unordered_map<std::string, ServiceUUID>;
  ProxyFactoryMap factoryMap_;      ///< Proxy factory map
  ProxyFactoryNameIDMap nameIdMap_; ///< {class name, Service UUID}
  static ProxyID proxyID_;          ///< Proxy starter ID

public:
  virtual ~ProxyCreatorImpl();

  // Create a proxy
  // @param serviceUUID service UUID
  // @param trans data stream
  // @param attached attach to a service instance or not
  // @param rpc Rpc
  // @return Proxy pointer
  virtual ProxyPtr getService(ServiceUUID serviceUUID, TransportPtr &trans,
                              bool attached, Rpc *rpc) override;
  // Create a proxy
  // @param className The name of class
  // @param trans data stream
  // @param attached attach to a service instance or not
  // @param rpc Rpc
  // @return Proxy pointer
  virtual ProxyPtr getService(const std::string &className, TransportPtr &trans,
                              bool attached, Rpc *rpc) override;
  // Registers a service factory
  // @param serviceUUID service UUID
  // @param factory service factory
  virtual void addFactory(ServiceUUID serviceUUID,
                          ProxyFactory *factory) override;
  // Make a new proxy instance ID
  virtual ProxyID createProxyID() override;

public:
  // Cleanup all stuff
  void cleanup();
};

// ProxyCall manager interface
class ProxyCallManager {
public:
  virtual ~ProxyCallManager() {}
  virtual void add(ProxyCallPtr call) = 0;
  // Find a proxy call
  // @param callID proxy call ID
  virtual ProxyCallPtr get(CallID callID) = 0;
  // Destroy proxy call
  // @param callID proxy call ID
  virtual void destroy(CallID callID) = 0;
  virtual void update(std::time_t ms = 0) = 0;
  // Invoked when proxy call timeout
  virtual void timeout(ProxyCallPtr &callPtr) = 0;
  // Retry
  virtual void retry(ProxyCallPtr callPtr) = 0;
  // Cleanup
  virtual void cleanup() = 0;
};

// Manages proxy's call
class ProxyCallManagerImpl : public ProxyCallManager {
  using CallMap = std::unordered_map<CallID, ProxyCallPtr>;
  CallMap callMap_;       ///< ProxyCall map
  Timer timer_;           ///< Timer manager
  CallID callID_{1};      ///< Proxy call starter ID
  RpcImpl *rpc_{nullptr}; ///< Rpc instance

public:
  ProxyCallManagerImpl(RpcImpl *rpc);
  virtual ~ProxyCallManagerImpl();
  // Add a proxy call
  // @param call proxy call instance
  virtual void add(ProxyCallPtr call) override;
  // Find a proxy call
  // @param callID proxy call ID
  virtual ProxyCallPtr get(CallID callID) override;
  // Destroy proxy call
  // @param callID proxy call ID
  virtual void destroy(CallID callID) override;
  // Main loop, invoked by Rpc::update()
  // @param ms current timestamp in millionseconds
  virtual void update(std::time_t ms = 0) override;
  // Invoked when proxy call timeout
  virtual void timeout(ProxyCallPtr &callPtr) override;
  // Returns timer manager
  TimerBase *getTimer();
  // Retry
  virtual void retry(ProxyCallPtr callPtr) override;
  // Cleanup
  virtual void cleanup() override;
  // Remove all proxy all by proxy ID
  void removeAll(ProxyID proxyID);
};

// Proxy method call
class ProxyCall {
  TimerHandle timerHandle_{nullptr}; ///< Timer handle
  int timeout_{RPC_TIMEOUT};         ///< RPC timeout
  CallID callID_{INVALID_CALL_ID};   ///< Proxy calL ID
  ErrorID errorID_{0};               ///< The error code for current call
  bool inCoro_{false};               ///< In coroutine
  int retry_{0};                     ///< Retry times
  char *data_{nullptr};              ///< Method's argument serialized data
  int dataSize_{0};                  ///< Argument serialized data size
  bool noExcept_{false};             ///< Throw exception or not
  std::string traceID_;              ///< Trace ID
  std::uint64_t spanID_{0};          ///< Trace span ID
  std::uint64_t parentSpanID_{0};    ///< Trace parent span ID
  std::time_t forkTimestamp_{0}; ///< The timestamp that ProxyCall was generated
#if defined(DEBUG) || defined(_DEBUG)
  std::string callName_; ///< For debug version call method name
#endif

protected:
  rpc::TransportPtr transport_{nullptr}; ///< Target service's transport
  Rpc *rpc_{nullptr};                    ///< Rpc instance

public:
  // dtor
  virtual ~ProxyCall();
  // Invoked by framework when return value arrived
  virtual void doRet(int bytes) = 0;
  // Returns proxy ID
  virtual ProxyID getProxyID() = 0;
  // Returns remote method name
  const std::string &getCallName();
  // Returns proxy call ID
  CallID getCallID();
  // Returns error ID
  ErrorID getErrorID();
  // Is coroutinized?
  bool isCoroutinized();
  // Setup coroutine flag
  void setCoro();
  // Returns retry times
  int getRetry();
  // Check retry count
  bool needRetry();
  // Setup retry times
  void setRetry(int retry);
  // Re-send
  void resend();
  // Start timer
  void startExecTimer();
  // Stop timer
  void stopExecTimer();
  // Setup data and size
  void setData(const char *data, int size);
  // Returns noexcept flag
  bool getNoExcept();
  // Returns Rpc
  Rpc *getRpc();
  void setTraceID(const std::string &traceID);
  const std::string &getTraceID();
  void setSpanID(std::uint64_t spanID);
  std::uint64_t getSpanID();
  void setParentSpanID(std::uint64_t parentSpanID);
  std::uint64_t getParentSpanID();
  rpc::TransportPtr getTransport();
  // Returns The timestamp that ProxyCall was generated
  std::time_t getForkTimestamp();
  void setForkTimestamp(std::time_t forkTimestamp);

protected:
  // Setup proxy call ID
  void setCallID(CallID callID);
  // Setup error ID
  void setErrorID(ErrorID errorID);
  // Setup remote method name
  void setCallName(const std::string &name);
  // Setup timer handler for current call
  void setTimerHandle(TimerHandle handle);
  // Setup RPC timeout in second
  void setTimeout(int timeout);
  // Returns RPC timeout
  int getTimeout();
  // Returns timer handle
  TimerHandle getTimerHandle();
  // Setup noexcept flag
  void setNoExcept(bool noExcept);

  friend class ProxyCallManagerImpl;
  friend class RpcImpl;
};

extern ProxyCreator &getGlobalProxyCreator();

#define ProxyCreatorRef (rpc::getGlobalProxyCreator())

template <typename T, typename TransT>
std::shared_ptr<T> getService(TransT transport, bool attached = true,
                              Rpc *rpc = nullptr) {
  if (!transport) {
    return nullptr;
  }
  return std::dynamic_pointer_cast<T>(
      rpc->getProxyCreator()->getService(T::uuid(), transport, attached, rpc));
}

template <typename T, typename TransT>
std::shared_ptr<T> getService(GlobalIndex globalIndex, TransT transport,
                              bool attached = true, Rpc *rpc = nullptr) {
  if (!transport) {
    return nullptr;
  }
  auto proxy = std::dynamic_pointer_cast<T>(
      rpc->getProxyCreator()->getService(T::uuid(), transport, attached, rpc));
  if (proxy) {
    proxy->setGlobalIndex(globalIndex);
  }
  return proxy;
}

template <typename T, typename TransT>
std::shared_ptr<T> createProxy(TransT transport, bool attached = true,
                               Rpc *rpc = nullptr) {

  if (!transport) {
    return nullptr;
  }
  return rpc::getService<T>(transport, attached, rpc);
}

template <typename T> bool registerServiceFactory() {
  getGlobalProxyCreator().addFactory(T::uuid(), new ProxyFactoryT<T>());
  return true;
}

extern void cleanup_all_proxy_stuff();

extern void buildHeader(RpcCallHeader &cheader, rpc::Proxy *proxy,
                        rpc::CallID callID, rpc::MethodID methodID,
                        std::uint32_t bodySize);

extern void buildHeader(RpcCallHeader &cheader, rpc::ServiceUUID service_uuid,
                        rpc::ServiceID service_id, rpc::CallID callID,
                        rpc::MethodID methodID, std::uint32_t bodySize);

extern void buildHeader(RpcProxyCallHeader &cheader, rpc::Proxy *proxy,
                        rpc::CallID callID, rpc::MethodID methodID,
                        std::uint32_t bodySize, bool oneway);

} // namespace rpc
